#include <gtest/gtest.h>

#include "common_components/common_runtime/src/heap/allocator/treap.h"


using namespace panda;

class TreapTest : public ::testing::Test {
protected:
    void SetUp() override {
        treap_.Init(1024); // 初始化足够大的内存块
    }

    void TearDown() override {
        treap_.Fini();
    }

    Treap treap_;
};

TEST_F(TreapTest, InitAndEmpty) {
    EXPECT_TRUE(treap_.Empty());
    EXPECT_EQ(treap_.GetTotalCount(), 0u);
}

TEST_F(TreapTest, InsertSingleNode) {
    bool result = treap_.MergeInsert(100, 10, true);
    EXPECT_TRUE(result);
    EXPECT_FALSE(treap_.Empty());
    EXPECT_EQ(treap_.GetTotalCount(), 10u);

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 100u);
    EXPECT_EQ(root->GetCount(), 10u);
}

TEST_F(TreapTest, MergeInsertAdjacentNodes) {
    // 插入第一个节点
    treap_.MergeInsert(100, 10, true);

    // 插入与前一个节点相邻的节点，应该合并
    bool result = treap_.MergeInsert(110, 5, true);
    EXPECT_TRUE(result);
    EXPECT_EQ(treap_.GetTotalCount(), 15u);

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 100u);
    EXPECT_EQ(root->GetCount(), 15u);
}

TEST_F(TreapTest, TakeUnitsFromTree) {
    // 插入一个大节点
    treap_.MergeInsert(200, 20, true);

    uint32_t idx = 0;
    bool result = treap_.TakeUnits(5, idx, true);
    EXPECT_TRUE(result);
    EXPECT_EQ(idx, 200u);

    const Treap::TreapNode* root = treap_.RootNode();
    ASSERT_NE(root, nullptr);
    EXPECT_EQ(root->GetIndex(), 205u);
    EXPECT_EQ(root->GetCount(), 15u);
    EXPECT_EQ(treap_.GetTotalCount(), 15u);
}

TEST_F(TreapTest, TakeUnitsFailOnSmallNode) {
    // 插入一个小节点
    treap_.MergeInsert(300, 3, true);

    uint32_t idx = 0;
    bool result = treap_.TakeUnits(5, idx, true);
    EXPECT_FALSE(result);
}

TEST_F(TreapTest, IteratorTraversal) {
    treap_.MergeInsert(400, 10, true);
    treap_.MergeInsert(410, 5, true);

    Treap::Iterator it(treap_);
    const Treap::TreapNode* node = it.Next();
    int count = 0;

    while (node != nullptr) {
        count++;
        node = it.Next();
    }

    EXPECT_EQ(count, 2); // 应该有两个节点
}

TEST_F(TreapTest, ZeroUnitInsertionFails) {
    bool result = treap_.MergeInsert(500, 0, true);
    EXPECT_FALSE(result);
}

TEST_F(TreapTest, EmptyTreeTakeFails) {
    uint32_t idx = 0;
    bool result = treap_.TakeUnits(5, idx, true);
    EXPECT_FALSE(result);
}